import networkx as nx
import matplotlib
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
from PIL import Image
import os
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
from collections import defaultdict
import utils.mylog as logging


class UTG:
    """
    TODO:记录控件每一步操作结果及UTG图
    """
    def __init__(self, device):
        self.logger = logging.getLogger(self.__class__.__name__)
        self.device = device
        self.states = {}    # 状态为节点
        self.events = []   # 操作为边


    def add_transition(self, event, old_state, new_state):
        self.add_node(old_state)
        self.add_node(new_state)
        # 检查边是否已存在，通过判断是否已有相同的起始和终止节点的边
        # existing_event = self._find_existing_edge(old_state, new_state)
        # if existing_event is None:
        event.build_edge(old_state, new_state)
        self.events.append(event)
        self.states[old_state.state_id].add_event(event)
        self.states[new_state.state_id].add_event(event)

    def add_node(self, state):
        if not state:
            return
        if state.state_id not in list(self.states.keys()):
            self.states[state.state_id] = state

    def _find_existing_edge(self, source_node, target_node):
        for edge in self.events:
            if edge.source_state == source_node and edge.target_state == target_node :
                return edge
        return None


    """
    TODO:绘制优化为实时绘制
    """
    def draw_graph(self):
        G = nx.DiGraph()

        # 将自定义图结构中的节点添加到networkx的图中
        for node in self.states:
            # TODO:hash_data调整为截图
            self.logger.info(f"add_node: {str(node)}")
            G.add_node(str(node))

        # 将自定义图结构中的边添加到networkx的图中
        for edge in self.events:
            self.logger.info(f"add_edge: from:{str(edge.source_status.state_id)} to:{str(edge.target_status.state_id)}")
            G.add_edge(str(edge.source_status.state_id), str(edge.target_status.state_id))

        # 布局节点，这里使用spring_layout布局算法，可尝试其他布局算法找到合适的展示效果
        pos = nx.spring_layout(G)

        # 绘制边
        nx.draw_networkx_edges(G, pos, edge_color='gray', arrows=True, arrowsize=20)

        # 节点
        # nx.draw_networkx_nodes(G, pos, node_size=500, node_color='lightblue')
        ax = plt.gca()
        for node in G.nodes():
            x, y = pos[node]
            img_path = self.states[node].screencap_dir
            if os.path.exists(img_path):
                img = plt.imread(img_path)
                im = OffsetImage(img, zoom=0.05)
                # 设置box_alignment参数，调整图片位置的对齐方式，使其更合理显示
                ab = AnnotationBbox(im, (x, y), xycoords='data', frameon=False, box_alignment=(0.5, 0.5), zorder=0)
                # 设置clip_on参数为False，避免图片裁剪边等元素
                ab.set_clip_on(False)
                ax.add_artist(ab)




        # 添加节点标签，这里简单地使用节点数据作为标签，可根据需求自定义更详细的标签
        # node_labels = {n: str(n) for n in G.nodes}
        # nx.draw_networkx_labels(G, pos, labels=node_labels, font_size=10)

        # 添加边标签（如果边有额外属性，比如权重等，可以在这里展示，当前示例边无其他属性则简单处理）
        # edge_labels = defaultdict(str)
        # nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)


        plt.axis('off')
        plt.show()

